#include "myDialect/toyDialect.h"
#include "myDialect/toyOps.h"
#include "rewrite/passes.h"

#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Support/TypeID.h"

#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/Support/Casting.h"
#include <memory>
#include <utility>
using namespace mlir;

namespace {

class PrintOpLowering : public OpRewritePattern<toy::PrintOp> {
    using OpRewritePattern<toy::PrintOp>::OpRewritePattern;
    
    LogicalResult matchAndRewrite(toy::PrintOp op, 
                                    PatternRewriter& rewriter) const final {
        auto location = op.getLoc();
        ModuleOp parentModule = op.getOperation()->getParentOfType<ModuleOp>();
        auto printfRef = this->getOrInsertPrintf(rewriter, parentModule);
        // Value input = op.getInput();
        // Value formatSpecifierCst = getOrCreateGlobalString(location, 
        //     rewriter, "frmt_spec", StringRef("%d \0", 4), parentModule);

        // rewriter.setInsertionPointAfter(op.getOperation());
        // rewriter.create<LLVM::CallOp>(location, getPrintfType(rewriter.getContext()),
        //                                     printfRef, ArrayRef<Value>({formatSpecifierCst, input}));
        rewriter.eraseOp(op.getOperation());
        return success();
    }

private:
    static LLVM::LLVMFunctionType getPrintfType(MLIRContext* context) {
        auto llvmI32Ty = IntegerType::get(context, 32);
        auto llvmPtrTy = LLVM::LLVMPointerType::get(context);
        auto llvmFnTy = LLVM::LLVMFunctionType::get(llvmI32Ty, llvmPtrTy, true);
        return llvmFnTy;
    }

    static FlatSymbolRefAttr getOrInsertPrintf(PatternRewriter& rewriter, ModuleOp module) {
        auto context = module.getContext();
        if(module.lookupSymbol<LLVM::LLVMFuncOp>("printf")){
            return SymbolRefAttr::get(context, "printf");
        }
        PatternRewriter::InsertionGuard insertGuard(rewriter);
        rewriter.setInsertionPointToStart(module.getBody());
        rewriter.create<LLVM::LLVMFuncOp>(module.getLoc(), "printf", getPrintfType(context));
        return SymbolRefAttr::get(context, "printf");
    }

    static Value getOrCreateGlobalString(Location loc, OpBuilder &builder,
                                        StringRef name, StringRef value,
                                        ModuleOp module) {
        // Create the global at the entry of the module.
        LLVM::GlobalOp global;
        if (!(global = module.lookupSymbol<LLVM::GlobalOp>(name))) {
        OpBuilder::InsertionGuard insertGuard(builder);
        builder.setInsertionPointToStart(module.getBody());
        auto type = LLVM::LLVMArrayType::get(
            IntegerType::get(builder.getContext(), 8), value.size());
        global = builder.create<LLVM::GlobalOp>(loc, type, /*isConstant=*/true,
                                                LLVM::Linkage::Internal, name,
                                                builder.getStringAttr(value),
                                                /*alignment=*/0);
        }

        // Get the pointer to the first character in the global string.
        Value globalPtr = builder.create<LLVM::AddressOfOp>(loc, global);
        Value cst0 = builder.create<LLVM::ConstantOp>(loc, builder.getI64Type(),
                                                    builder.getIndexAttr(0));
        return builder.create<LLVM::GEPOp>(
            loc, LLVM::LLVMPointerType::get(builder.getContext()), global.getType(),
            globalPtr, ArrayRef<Value>({cst0, cst0}));
    }

};
}

namespace {
struct ToyToLLVMLoweringPass : 
    public PassWrapper<ToyToLLVMLoweringPass, OperationPass<ModuleOp>> {
    MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ToyToLLVMLoweringPass)

    void getDependentDialects(DialectRegistry& registry) const override {
        registry.insert<LLVM::LLVMDialect>();
    }
    void runOnOperation() final;
};
}

void ToyToLLVMLoweringPass::runOnOperation() {
    LLVMConversionTarget target(getContext());
    target.addLegalOp<ModuleOp>();
    LLVMTypeConverter typeConverter(&getContext());

    RewritePatternSet patterns(&getContext());
    mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);

    patterns.add<PrintOpLowering>(&getContext());

    auto module = getOperation();
    if (failed(applyFullConversion(module, target, std::move(patterns))))
        signalPassFailure();
}

std::unique_ptr<mlir::Pass> mlir::toy::createLowerToLLVMPass() {
  return std::make_unique<ToyToLLVMLoweringPass>();
}